﻿﻿﻿﻿﻿﻿﻿﻿using System.Security.Cryptography;
using System.Text;
using System.Collections.Concurrent;
using System.Buffers;
using DeviceCommons.Validation;

namespace DeviceCommons.Security
{
    /// <summary>
    /// 高性能AES加密器，支持密钥缓存和快速模式
    /// </summary>
    public class AesEncryptor
    {
        private static readonly int DefaultKeySize = 256;
        private static readonly int DefaultDerivationIterations = 60000;
        private static readonly int FastModeIterations = 1000; // 快速模式：1000次迭代
        private static readonly int DefaultSaltSize = 16;
        private static readonly int DefaultIvSize = 16;
        
        // 密钥缓存：避免重复计算相同密码的PBKDF2
        private readonly ConcurrentDictionary<string, byte[]> _keyCache;
        private readonly object _cacheLock = new object();
        private const int MaxCacheSize = 100; // 最大缓存100个密钥
        
        private readonly int keySize;
        private readonly int derivationIterations;
        private readonly int saltSize;
        private readonly int ivSize;
        private readonly bool enableKeyCache;
        
        /// <summary>
        /// 创建默认AES加密器（安全模式，60000次迭代）
        /// </summary>
        public AesEncryptor()
            : this(DefaultKeySize, DefaultDerivationIterations, DefaultSaltSize, DefaultIvSize, false)
        {
            _keyCache = new ConcurrentDictionary<string, byte[]>();
        }
        
        /// <summary>
        /// 创建快速AES加密器（性能模式，1000次迭代，启用缓存）
        /// </summary>
        /// <param name="fastMode">是否启用快速模式</param>
        /// <param name="enableCache">是否启用密钥缓存</param>
        public AesEncryptor(bool fastMode, bool enableCache = true)
            : this(DefaultKeySize, 
                   fastMode ? FastModeIterations : DefaultDerivationIterations, 
                   DefaultSaltSize, DefaultIvSize, enableCache)
        {
            _keyCache = enableCache ? new ConcurrentDictionary<string, byte[]>() : null;
        }

        public AesEncryptor(int keySize, int derivationIterations, int saltSize, int ivSize, bool enableKeyCache = false)
        {
            if (keySize != 128 && keySize != 192 && keySize != 256)
                throw new ArgumentOutOfRangeException(nameof(keySize), "Key size must be 128, 192, or 256 bits.");
            if (derivationIterations <= 0)
                throw new ArgumentOutOfRangeException(nameof(derivationIterations), "Derivation iterations must be greater than zero.");
            if (saltSize <= 0)
                throw new ArgumentOutOfRangeException(nameof(saltSize), "Salt size must be greater than zero.");
            if (ivSize <= 0)
                throw new ArgumentOutOfRangeException(nameof(ivSize), "IV size must be greater than zero.");

            this.keySize = keySize;
            this.derivationIterations = derivationIterations;
            this.saltSize = saltSize;
            this.ivSize = ivSize;
            this.enableKeyCache = enableKeyCache;
            this._keyCache = enableKeyCache ? new ConcurrentDictionary<string, byte[]>() : null;
        }

        /// <summary>
        /// 高性能加密方法，支持密钥缓存和内存优化
        /// </summary>
        public string Encrypt(string plainText, string passPhrase)
        {
            if (string.IsNullOrEmpty(plainText)) throw new ArgumentNullException(nameof(plainText), "Plain text cannot be empty");
            if (string.IsNullOrEmpty(passPhrase)) throw new ArgumentNullException(nameof(passPhrase), "Passphrase cannot be empty");

            // 使用内存池减少内存分配
            var saltBuffer = ArrayPool<byte>.Shared.Rent(saltSize);
            var ivBuffer = ArrayPool<byte>.Shared.Rent(ivSize);
            
            try
            {
                var salt = saltBuffer.AsSpan(0, saltSize);
                var iv = ivBuffer.AsSpan(0, ivSize);
                
                // 使用更高效的随机数生成
                RandomNumberGenerator.Fill(salt);
                RandomNumberGenerator.Fill(iv);

                var keyBytes = DeriveKeyBytesOptimized(passPhrase, salt.ToArray());

                using var symmetricKey = CreateAesSymmetricKey(keyBytes, iv.ToArray());
                using var memoryStream = new MemoryStream();

                // 直接写入字节数据，避免不必要的复制
                memoryStream.Write(salt);
                memoryStream.Write(iv);

                using (var cryptoStream = new CryptoStream(memoryStream,
                                                          symmetricKey.CreateEncryptor(),
                                                          CryptoStreamMode.Write))
                {
                    var plainTextBytes = Encoding.UTF8.GetBytes(plainText);
                    cryptoStream.Write(plainTextBytes, 0, plainTextBytes.Length);
                    cryptoStream.FlushFinalBlock();
                }

                return Convert.ToBase64String(memoryStream.ToArray());
            }
            finally
            {
                // 归还内存池缓冲区
                ArrayPool<byte>.Shared.Return(saltBuffer);
                ArrayPool<byte>.Shared.Return(ivBuffer);
            }
        }

        /// <summary>
        /// 高性能解密方法，支持密钥缓存和内存优化
        /// </summary>
        public string Decrypt(string cipherText, string passPhrase)
        {
            if (string.IsNullOrEmpty(cipherText)) throw new ArgumentNullException(nameof(cipherText), "Cipher text cannot be empty");
            if (string.IsNullOrEmpty(passPhrase)) throw new ArgumentNullException(nameof(passPhrase), "Passphrase cannot be empty");

            try
            {
                var cipherTextBytesWithSaltAndIv = Convert.FromBase64String(cipherText);

                if (cipherTextBytesWithSaltAndIv.Length < saltSize + ivSize)
                    throw new DeviceMessageValidationException(
                        ValidationErrorType.InvalidCipherTextFormat,
                        $"密文格式无效，数据长度不足。期望至少 {saltSize + ivSize} 字节，但实际获得 {cipherTextBytesWithSaltAndIv.Length} 字节",
                        nameof(cipherText),
                        $"至少 {saltSize + ivSize} 字节",
                        $"{cipherTextBytesWithSaltAndIv.Length} 字节");

                // 使用Span提高性能，避免不必要的数组复制
                var dataSpan = cipherTextBytesWithSaltAndIv.AsSpan();
                var salt = dataSpan.Slice(0, saltSize).ToArray();
                var iv = dataSpan.Slice(saltSize, ivSize).ToArray();
                var cipherTextBytes = dataSpan.Slice(saltSize + ivSize).ToArray();

                var keyBytes = DeriveKeyBytesOptimized(passPhrase, salt);

                using var symmetricKey = CreateAesSymmetricKey(keyBytes, iv);
                using var memoryStream = new MemoryStream(cipherTextBytes);
                using var cryptoStream = new CryptoStream(memoryStream, symmetricKey.CreateDecryptor(), CryptoStreamMode.Read);
                
                // 直接读取字节数据，避免StreamReader的开销
                var decryptedBytes = new List<byte>();
                int b;
                while ((b = cryptoStream.ReadByte()) != -1)
                {
                    decryptedBytes.Add((byte)b);
                }
                
                return Encoding.UTF8.GetString(decryptedBytes.ToArray());
            }
            catch (CryptographicException ex)
            {
                throw new InvalidOperationException("Invalid passphrase or corrupted data.", ex);
            }
            catch (FormatException ex)
            {
                throw new DeviceMessageValidationException(
                    ValidationErrorType.InvalidBase64Format,
                    $"密文Base64格式无效: {cipherText}",
                    ex,
                    nameof(cipherText));
            }
        }

        public async ValueTask<string> EncryptAsync(string plainText, string passPhrase, CancellationToken cancellationToken = default)
        {
            if (string.IsNullOrEmpty(plainText)) throw new ArgumentNullException(nameof(plainText));
            if (string.IsNullOrEmpty(passPhrase)) throw new ArgumentNullException(nameof(passPhrase));

            var salt = GenerateRandomBytes(saltSize);
            var iv = GenerateRandomBytes(ivSize);
            var keyBytes = DeriveKeyBytes(passPhrase, salt);

            using var symmetricKey = CreateAesSymmetricKey(keyBytes, iv);
            using var memoryStream = new MemoryStream();
            await memoryStream.WriteAsync(salt, cancellationToken).ConfigureAwait(false);
            await memoryStream.WriteAsync(iv, cancellationToken).ConfigureAwait(false);

            using (var cryptoStream = new CryptoStream(memoryStream, symmetricKey.CreateEncryptor(), CryptoStreamMode.Write))
            using (var writer = new StreamWriter(cryptoStream, Encoding.UTF8))
            {
                await writer.WriteAsync(plainText).ConfigureAwait(false);
                await writer.FlushAsync().ConfigureAwait(false);
            }

            return Convert.ToBase64String(memoryStream.ToArray());
        }

        public async ValueTask<string> DecryptAsync(string cipherText, string passPhrase, CancellationToken cancellationToken = default)
        {
            if (string.IsNullOrEmpty(cipherText)) throw new ArgumentNullException(nameof(cipherText));
            if (string.IsNullOrEmpty(passPhrase)) throw new ArgumentNullException(nameof(passPhrase));

            try
            {
                var cipherTextBytesWithSaltAndIv = Convert.FromBase64String(cipherText);

                if (cipherTextBytesWithSaltAndIv.Length < saltSize + ivSize)
                    throw new DeviceMessageValidationException(
                        ValidationErrorType.InvalidCipherTextFormat,
                        "密文格式无效，数据长度不足",
                        nameof(cipherText),
                        $"至少 {saltSize + ivSize} 字节",
                        $"{cipherTextBytesWithSaltAndIv.Length} 字节");

                var salt = cipherTextBytesWithSaltAndIv.Take(saltSize).ToArray();
                var iv = cipherTextBytesWithSaltAndIv.Skip(saltSize).Take(ivSize).ToArray();
                var cipherTextBytes = cipherTextBytesWithSaltAndIv.Skip(saltSize + ivSize).ToArray();

                var keyBytes = DeriveKeyBytes(passPhrase, salt);

                using var symmetricKey = CreateAesSymmetricKey(keyBytes, iv);
                using var memoryStream = new MemoryStream(cipherTextBytes);
                using var cryptoStream = new CryptoStream(memoryStream, symmetricKey.CreateDecryptor(), CryptoStreamMode.Read);
                using var reader = new StreamReader(cryptoStream, Encoding.UTF8);

                return await reader.ReadToEndAsync().ConfigureAwait(false);
            }
            catch (Exception ex)
            {
                throw new InvalidOperationException("Decryption failed", ex);
            }
        }

        /// <summary>
        /// 优化的密钥派生方法，支持缓存
        /// </summary>
        private byte[] DeriveKeyBytesOptimized(string passPhrase, byte[] salt)
        {
            if (!enableKeyCache)
            {
                // 不启用缓存时直接计算
                return DeriveKeyBytesDirect(passPhrase, salt);
            }

            // 使用密码+盐值作为缓存键
            var cacheKey = $"{passPhrase}:{Convert.ToBase64String(salt)}";
            
            if (_keyCache.TryGetValue(cacheKey, out var cachedKey))
            {
                return cachedKey;
            }

            // 计算新密钥
            var keyBytes = DeriveKeyBytesDirect(passPhrase, salt);

            // 缓存管理：防止缓存过大
            lock (_cacheLock)
            {
                if (_keyCache.Count >= MaxCacheSize)
                {
                    // 清理一半缓存
                    var keysToRemove = _keyCache.Keys.Take(MaxCacheSize / 2).ToArray();
                    foreach (var key in keysToRemove)
                    {
                        _keyCache.TryRemove(key, out _);
                    }
                }
                
                _keyCache.TryAdd(cacheKey, keyBytes);
            }

            return keyBytes;
        }
        
        /// <summary>
        /// 直接计算密钥派生（无缓存）
        /// </summary>
        private byte[] DeriveKeyBytesDirect(string passPhrase, byte[] salt)
        {
            // 使用更高效的PBKDF2实现
            using var pbkdf2 = new Rfc2898DeriveBytes(passPhrase, salt, derivationIterations, HashAlgorithmName.SHA256);
            return pbkdf2.GetBytes(keySize / 8);
        }

        private byte[] DeriveKeyBytes(string passPhrase, byte[] salt)
        {
            // 保持向后兼容
            return DeriveKeyBytesOptimized(passPhrase, salt);
        }

        private Aes CreateAesSymmetricKey(byte[] keyBytes, byte[] iv)
        {
            var aes = Aes.Create();
            aes.Mode = CipherMode.CBC;
            aes.Padding = PaddingMode.PKCS7;
            aes.Key = keyBytes;
            aes.IV = iv;
            return aes;
        }

        /// <summary>
        /// 生成随机字节（向后兼容）
        /// </summary>
        private byte[] GenerateRandomBytes(int size)
        {
            var randomBytes = new byte[size];
            RandomNumberGenerator.Fill(randomBytes);
            return randomBytes;
        }
        
        /// <summary>
        /// 清空当前实例的密钥缓存
        /// </summary>
        public void ClearKeyCache()
        {
            if (_keyCache != null)
            {
                lock (_cacheLock)
                {
                    _keyCache.Clear();
                }
            }
        }
        
        /// <summary>
        /// 获取当前实例的缓存统计信息
        /// </summary>
        public (int Count, int MaxSize) GetCacheStats()
        {
            if (_keyCache == null)
                return (0, MaxCacheSize);
                
            return (_keyCache.Count, MaxCacheSize);
        }
        
        /// <summary>
        /// 创建性能模式的AES加密器（用于性能测试）
        /// </summary>
        /// <returns>高性能加密器实例</returns>
        public static AesEncryptor CreateFastMode()
        {
            return new AesEncryptor(true, true); // 快速模式 + 启用缓存
        }
        
        /// <summary>
        /// 创建安全模式的AES加密器（用于生产环境）
        /// </summary>
        /// <returns>安全加密器实例</returns>
        public static AesEncryptor CreateSecureMode()
        {
            return new AesEncryptor(false, false); // 安全模式 + 禁用缓存
        }
    }
}
